import pandas as pd
from faker import Faker
import numpy as np
import random
import time
from typing import Dict, List, Any, Optional, Union
import os

class DataGenerator:
    """数据生成器类，用于生成各种类型的测试数据"""
    
    def __init__(self, locale: str = 'zh_CN'):
        """初始化数据生成器
        
        Args:
            locale: 地区设置，默认为中文
        """
        self.faker = Faker(locale)
        Faker.seed(42)  # 设置随机种子，保证结果可复现
        random.seed(42)
        np.random.seed(42)
    
    def generate_user_data(self, count: int, output_file: Optional[str] = None) -> pd.DataFrame:
        """生成用户数据
        
        Args:
            count: 生成数据的数量
            output_file: 输出文件路径，可选
        
        Returns:
            包含用户数据的DataFrame
        """
        start_time = time.time()
        print(f"开始生成{count}条用户数据...")
        
        # 使用列表推导式高效生成数据
        data = [{
            'user_id': i + 1,
            'username': self.faker.user_name(),
            'full_name': self.faker.name(),
            'email': self.faker.email(),
            'phone': self.faker.phone_number()[:11],  # 确保手机号为11位
            'address': self.faker.address(),
            'birthday': self.faker.date_of_birth(minimum_age=18, maximum_age=70),
            'gender': random.choice(['男', '女']),
            'registration_date': self.faker.date_time_between(start_date='-3y', end_date='now'),
            'last_login': self.faker.date_time_between(start_date='-1y', end_date='now'),
            'status': random.choice(['活跃', '休眠', '冻结']),
            'credit_score': random.randint(300, 850)
        } for i in range(count)]
        
        df = pd.DataFrame(data)
        
        # 转换日期格式为字符串
        df['birthday'] = pd.to_datetime(df['birthday'])
        df['registration_date'] = pd.to_datetime(df['registration_date'])
        df['last_login'] = pd.to_datetime(df['last_login'])
        
        if output_file:
            self._save_to_file(df, output_file)
        
        end_time = time.time()
        print(f"用户数据生成完成，耗时: {end_time - start_time:.2f}秒")
        return df
    
    def generate_order_data(self, count: int, user_ids: List[int], output_file: Optional[str] = None) -> pd.DataFrame:
        """生成订单数据
        
        Args:
            count: 生成数据的数量
            user_ids: 用户ID列表
            output_file: 输出文件路径，可选
        
        Returns:
            包含订单数据的DataFrame
        """
        start_time = time.time()
        print(f"开始生成{count}条订单数据...")
        
        # 商品列表
        products = [
            {'name': '智能手机A', 'price': 3999.00, 'category': '电子产品'},
            {'name': '笔记本电脑B', 'price': 5999.00, 'category': '电子产品'},
            {'name': '无线耳机C', 'price': 799.00, 'category': '电子产品'},
            {'name': '智能手表D', 'price': 1299.00, 'category': '电子产品'},
            {'name': '运动鞋E', 'price': 499.00, 'category': '服装鞋帽'},
            {'name': '休闲裤F', 'price': 299.00, 'category': '服装鞋帽'},
            {'name': '护肤套装G', 'price': 599.00, 'category': '美妆个护'},
            {'name': '厨房料理机H', 'price': 899.00, 'category': '家居家电'},
        ]
        
        # 支付方式
        payment_methods = ['支付宝', '微信支付', '银联', '信用卡']
        
        # 订单状态
        order_statuses = ['待支付', '已支付', '待发货', '已发货', '已完成', '已取消', '已退款']
        
        # 批量生成数据
        batch_size = 100000  # 每批次处理的记录数，防止内存溢出
        all_data = []
        
        for batch_start in range(0, count, batch_size):
            batch_count = min(batch_size, count - batch_start)
            batch_data = [{
                'order_id': f"ORD{batch_start + i + 1:08d}",
                'user_id': random.choice(user_ids),
                'product_info': random.choice(products),
                'order_amount': random.randint(100, 10000) + round(random.random(), 2),
                'payment_method': random.choice(payment_methods),
                'order_status': random.choice(order_statuses),
                'order_time': self.faker.date_time_between(start_date='-1y', end_date='now'),
                'payment_time': None,
                'shipping_address': self.faker.address(),
                'shipping_phone': self.faker.phone_number()[:11],
                'tracking_number': f"TRK{random.randint(1000000000, 9999999999)}"
            } for i in range(batch_count)]
            
            # 为已支付的订单设置支付时间
            for item in batch_data:
                if item['order_status'] in ['已支付', '待发货', '已发货', '已完成', '已退款']:
                    item['payment_time'] = self.faker.date_time_between(start_date=item['order_time'], end_date='now')
                else:
                    item['payment_time'] = None
            
            all_data.extend(batch_data)
        
        df = pd.DataFrame(all_data)
        
        # 展开product_info字典
        df = pd.concat([
            df.drop(['product_info'], axis=1),
            df['product_info'].apply(pd.Series)
        ], axis=1)
        
        # 转换日期格式
        df['order_time'] = pd.to_datetime(df['order_time'])
        df['payment_time'] = pd.to_datetime(df['payment_time'])
        
        if output_file:
            self._save_to_file(df, output_file)
        
        end_time = time.time()
        print(f"订单数据生成完成，耗时: {end_time - start_time:.2f}秒")
        return df
    
    def generate_transaction_data(self, count: int, output_file: Optional[str] = None) -> pd.DataFrame:
        """生成交易数据
        
        Args:
            count: 生成数据的数量
            output_file: 输出文件路径，可选
        
        Returns:
            包含交易数据的DataFrame
        """
        start_time = time.time()
        print(f"开始生成{count}条交易数据...")
        
        # 交易类型
        transaction_types = ['充值', '提现', '转账', '消费', '退款', '缴费']
        
        # 交易状态
        transaction_statuses = ['成功', '失败', '处理中']
        
        # 批量生成数据
        batch_size = 100000
        all_data = []
        
        for batch_start in range(0, count, batch_size):
            batch_count = min(batch_size, count - batch_start)
            batch_data = [{
                'transaction_id': f"TXN{batch_start + i + 1:010d}",
                'user_id': random.randint(1, 10000),
                'transaction_type': random.choice(transaction_types),
                'amount': random.randint(1, 100000) + round(random.random(), 2),
                'currency': 'CNY',
                'transaction_status': random.choice(transaction_statuses),
                'transaction_time': self.faker.date_time_between(start_date='-1y', end_date='now'),
                'source_account': f"{random.randint(1000, 9999)} **** **** {random.randint(1000, 9999)}",
                'target_account': f"{random.randint(1000, 9999)} **** **** {random.randint(1000, 9999)}",
                'channel': random.choice(['APP', 'WEB', 'WAP', 'POS', 'ATM']),
                'ip_address': self.faker.ipv4(),
                'device_id': f"DEV{random.randint(1000000000, 9999999999)}"
            } for i in range(batch_count)]
            
            all_data.extend(batch_data)
        
        df = pd.DataFrame(all_data)
        df['transaction_time'] = pd.to_datetime(df['transaction_time'])
        
        if output_file:
            self._save_to_file(df, output_file)
        
        end_time = time.time()
        print(f"交易数据生成完成，耗时: {end_time - start_time:.2f}秒")
        return df
    
    def generate_custom_data(self, count: int, schema: Dict[str, Any], output_file: Optional[str] = None) -> pd.DataFrame:
        """根据自定义模式生成数据
        
        Args:
            count: 生成数据的数量
            schema: 数据模式，定义每个字段的生成方式
            output_file: 输出文件路径，可选
        
        Returns:
            包含自定义数据的DataFrame
        """
        start_time = time.time()
        print(f"开始生成{count}条自定义数据...")
        
        # 批量生成数据
        batch_size = 100000
        all_data = []
        
        for batch_start in range(0, count, batch_size):
            batch_count = min(batch_size, count - batch_start)
            batch_data = []
            
            for i in range(batch_count):
                row = {}
                for field_name, field_config in schema.items():
                    if isinstance(field_config, dict):
                        # 支持配置生成方式
                        if 'type' in field_config:
                            if field_config['type'] == 'int':
                                min_val = field_config.get('min', 0)
                                max_val = field_config.get('max', 100)
                                row[field_name] = random.randint(min_val, max_val)
                            elif field_config['type'] == 'float':
                                min_val = field_config.get('min', 0.0)
                                max_val = field_config.get('max', 1.0)
                                row[field_name] = round(random.uniform(min_val, max_val), 2)
                            elif field_config['type'] == 'string':
                                length = field_config.get('length', 10)
                                row[field_name] = self.faker.pystr(min_chars=length, max_chars=length)
                            elif field_config['type'] == 'date':
                                start_date = field_config.get('start_date', '-1y')
                                end_date = field_config.get('end_date', 'now')
                                row[field_name] = self.faker.date_time_between(start_date=start_date, end_date=end_date)
                            elif field_config['type'] == 'choice':
                                choices = field_config.get('choices', [])
                                row[field_name] = random.choice(choices)
                    elif callable(field_config):
                        # 支持自定义生成函数
                        row[field_name] = field_config(self.faker, i)
                    else:
                        # 直接使用固定值
                        row[field_name] = field_config
                
                batch_data.append(row)
            
            all_data.extend(batch_data)
        
        df = pd.DataFrame(all_data)
        
        if output_file:
            self._save_to_file(df, output_file)
        
        end_time = time.time()
        print(f"自定义数据生成完成，耗时: {end_time - start_time:.2f}秒")
        return df
    
    def _save_to_file(self, df: pd.DataFrame, file_path: str) -> None:
        """将DataFrame保存到文件
        
        Args:
            df: 要保存的DataFrame
            file_path: 文件路径
        """
        # 确保目录存在
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        
        # 根据文件扩展名选择保存方式
        file_ext = os.path.splitext(file_path)[1].lower()
        
        if file_ext == '.csv':
            # 对于大型数据，使用分块写入
            if len(df) > 1000000:
                df.to_csv(file_path, index=False, chunksize=100000)
            else:
                df.to_csv(file_path, index=False)
        elif file_ext == '.parquet':
            # Parquet格式更适合大数据存储和分析
            df.to_parquet(file_path, index=False)
        elif file_ext == '.xlsx':
            # Excel格式适用于小数据量
            if len(df) > 100000:
                print("警告：数据量较大，不建议使用Excel格式保存")
            df.to_excel(file_path, index=False)
        else:
            # 默认使用CSV格式
            df.to_csv(file_path, index=False)
        
        print(f"数据已保存至: {file_path}")

# 示例用法
if __name__ == "__main__":
    # 创建数据生成器实例
    generator = DataGenerator()

    # 生成1000条交易数据
    # transactions_df = generator.generate_transaction_data(1000, './data/transactions.csv')
    # print(f"交易数据样例:\n{transactions_df.head()}")
    
    # 生成自定义数据
    # custom_schema = {
    #     'id': lambda faker, i: i + 1,
    #     'name': lambda faker, i: faker.name(),
    #     'age': {'type': 'int', 'min': 18, 'max': 65},
    #     'department': {'type': 'choice', 'choices': ['销售', '研发', '人事', '财务', '市场']},
    #     'salary': {'type': 'float', 'min': 3000, 'max': 20000},
    #     'hire_date': {'type': 'date', 'start_date': '-5y', 'end_date': 'now'},
    #     'active': True
    # }
    # custom_df = generator.generate_custom_data(100, custom_schema, './data/employees.csv')
    # print(f"自定义数据样例:\n{custom_df.head()}")